


import torch
from prior_depth_anything import PriorDepthAnything

device = "cuda:0" if torch.cuda.is_available() else "cpu"

ckpt_dir ='checkpoints'


# priorda = PriorDepthAnything(device=device, fmde_dir=ckpt_dir, ckpt_dir=ckpt_dir,cmde_dir = ckpt_dir)
priorda = PriorDepthAnything(device=device, fmde_dir=ckpt_dir, ckpt_dir=ckpt_dir,cmde_dir = ckpt_dir, frozen_model_size = 'vitl', conditioned_model_size='vitb')

image_path = 'assets/sample-2/rgb.jpg'
prior_path = 'assets/sample-2/prior_depth.png'


# image_path = 'assets/sample-6/rgb.npy'
# prior_path = 'assets/sample-6/prior_depth.npy'

output = priorda.infer_one_sample(image=image_path, prior=prior_path, visualize=True)
